import pygame
import random
import sys

# 初始化 Pygame
pygame.init()

# 定义颜色
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BACKGROUND = (40, 44, 52)  # 深色背景
GRID_COLOR = (48, 52, 60)  # 网格颜色
SNAKE_COLOR = (152, 195, 121)  # 柔和的绿色
FOOD_COLOR = (224, 108, 117)  # 柔和的红色
BORDER_COLOR = (86, 92, 100)  # 深灰色边框
SCORE_COLOR = (229, 192, 123)  # 金黄色分数
MENU_BG = (30, 34, 42)    # 菜单背景色
BUTTON_COLOR = (86, 92, 100)  # 按钮颜色
BUTTON_HOVER = (96, 102, 110)  # 按钮悬停颜色
TITLE_COLOR = (229, 192, 123)  # 标题颜色
GOLD_FOOD_COLOR = (255, 215, 0)  # 金色食物
SPECIAL_SCORE_COLOR = (255, 215, 0)  # 特殊分数颜色
PAUSE_BG = (0, 0, 0, 128)  # 半透明黑色

# 设置游戏窗口
WINDOW_WIDTH = 800
WINDOW_HEIGHT = 600
BLOCK_SIZE = 20
GAME_SPEED = 10

# 创建游戏窗口
screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption('贪吃蛇游戏')

# 蛇类
class Snake:
    def __init__(self):
        self.length = 1
        self.positions = [(WINDOW_WIDTH//2, WINDOW_HEIGHT//2)]
        self.direction = random.choice([UP, DOWN, LEFT, RIGHT])
        self.color = SNAKE_COLOR
        self.score = 0

    def get_head_position(self):
        return self.positions[0]

    def update(self):
        cur = self.get_head_position()
        x, y = self.direction
        new = (cur[0] + (x*BLOCK_SIZE), cur[1] + (y*BLOCK_SIZE))
        
        # 检查是否撞到边界
        if (new[0] < BLOCK_SIZE or new[0] >= WINDOW_WIDTH - BLOCK_SIZE or
            new[1] < BLOCK_SIZE or new[1] >= WINDOW_HEIGHT - BLOCK_SIZE):
            return False
        
        # 检查是否撞到自己
        if new in self.positions[3:]:
            return False
            
        self.positions.insert(0, new)
        if len(self.positions) > self.length:
            self.positions.pop()
        return True

    def reset(self):
        # 在重置之前检查是否需要更新最高分
        high_score = load_high_score()
        if self.score > high_score:
            save_high_score(self.score)
            
        self.length = 1
        self.positions = [(WINDOW_WIDTH//2, WINDOW_HEIGHT//2)]
        self.direction = random.choice([UP, DOWN, LEFT, RIGHT])
        self.score = 0

    def render(self):
        for i, p in enumerate(self.positions):
            # 蛇身绘制为圆角矩形
            if i == 0:  # 蛇头
                pygame.draw.rect(screen, self.color, (p[0], p[1], BLOCK_SIZE, BLOCK_SIZE), border_radius=7)
                # 添加眼睛
                eye_color = BLACK
                if self.direction == RIGHT:
                    pygame.draw.circle(screen, eye_color, (p[0] + BLOCK_SIZE - 5, p[1] + 5), 2)
                    pygame.draw.circle(screen, eye_color, (p[0] + BLOCK_SIZE - 5, p[1] + BLOCK_SIZE - 5), 2)
                elif self.direction == LEFT:
                    pygame.draw.circle(screen, eye_color, (p[0] + 5, p[1] + 5), 2)
                    pygame.draw.circle(screen, eye_color, (p[0] + 5, p[1] + BLOCK_SIZE - 5), 2)
                elif self.direction == UP:
                    pygame.draw.circle(screen, eye_color, (p[0] + 5, p[1] + 5), 2)
                    pygame.draw.circle(screen, eye_color, (p[0] + BLOCK_SIZE - 5, p[1] + 5), 2)
                elif self.direction == DOWN:
                    pygame.draw.circle(screen, eye_color, (p[0] + 5, p[1] + BLOCK_SIZE - 5), 2)
                    pygame.draw.circle(screen, eye_color, (p[0] + BLOCK_SIZE - 5, p[1] + BLOCK_SIZE - 5), 2)
            else:  # 蛇身
                pygame.draw.rect(screen, self.color, (p[0], p[1], BLOCK_SIZE, BLOCK_SIZE), border_radius=5)

# 食物类
class Food:
    def __init__(self):
        self.position = (0, 0)
        self.color = FOOD_COLOR
        self.is_special = False
        self.randomize_position()

    def randomize_position(self):
        self.position = (random.randint(0, (WINDOW_WIDTH-BLOCK_SIZE)//BLOCK_SIZE) * BLOCK_SIZE,
                        random.randint(0, (WINDOW_HEIGHT-BLOCK_SIZE)//BLOCK_SIZE) * BLOCK_SIZE)
        # 10%概率生成特殊食物
        self.is_special = random.random() < 0.1
        self.color = GOLD_FOOD_COLOR if self.is_special else FOOD_COLOR

    def render(self):
        center = (self.position[0] + BLOCK_SIZE // 2, self.position[1] + BLOCK_SIZE // 2)
        pygame.draw.circle(screen, self.color, center, BLOCK_SIZE // 2 - 2)
        if self.is_special:
            # 特殊食物添加闪光效果
            pygame.draw.circle(screen, WHITE, center, BLOCK_SIZE // 3, 1)

# 定义方向
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

def draw_grid():
    for x in range(BLOCK_SIZE, WINDOW_WIDTH - BLOCK_SIZE, BLOCK_SIZE):
        pygame.draw.line(screen, GRID_COLOR, (x, BLOCK_SIZE), (x, WINDOW_HEIGHT - BLOCK_SIZE))
    for y in range(BLOCK_SIZE, WINDOW_HEIGHT - BLOCK_SIZE, BLOCK_SIZE):
        pygame.draw.line(screen, GRID_COLOR, (BLOCK_SIZE, y), (WINDOW_WIDTH - BLOCK_SIZE, y))

def draw_border():
    # 绘制稍微粗一点的边框
    border_width = 2
    pygame.draw.rect(screen, BORDER_COLOR, (0, 0, WINDOW_WIDTH, BLOCK_SIZE))
    pygame.draw.rect(screen, BORDER_COLOR, (0, WINDOW_HEIGHT-BLOCK_SIZE, WINDOW_WIDTH, BLOCK_SIZE))
    pygame.draw.rect(screen, BORDER_COLOR, (0, 0, BLOCK_SIZE, WINDOW_HEIGHT))
    pygame.draw.rect(screen, BORDER_COLOR, (WINDOW_WIDTH-BLOCK_SIZE, 0, BLOCK_SIZE, WINDOW_HEIGHT))

# 在主要游戏逻辑之前添加按钮类
class Button:
    def __init__(self, x, y, width, height, text, font_size=36):
        self.rect = pygame.Rect(x, y, width, height)
        self.text = text
        self.font_size = font_size
        self.color = BUTTON_COLOR
        self.is_hovered = False
        
    def draw(self, surface):
        color = BUTTON_HOVER if self.is_hovered else BUTTON_COLOR
        pygame.draw.rect(surface, color, self.rect, border_radius=10)
        
        # 渲染按钮文字
        try:
            font = pygame.font.Font('arial.ttf', self.font_size)
        except:
            font = pygame.font.Font(None, self.font_size)
            
        text_surface = font.render(self.text, True, WHITE)
        text_rect = text_surface.get_rect(center=self.rect.center)
        surface.blit(text_surface, text_rect)
    
    def handle_event(self, event):
        if event.type == pygame.MOUSEMOTION:
            self.is_hovered = self.rect.collidepoint(event.pos)
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if self.is_hovered:
                return True
        return False

# 添加保存和加载最高分的函数
def save_high_score(score):
    try:
        with open('high_score.txt', 'w') as f:
            f.write(str(score))
    except:
        pass

def load_high_score():
    try:
        with open('high_score.txt', 'r') as f:
            return int(f.read())
    except:
        return 0

# 添加菜单界面函数
def show_menu():
    high_score = load_high_score()
    # 创建多个按钮
    start_button = Button(WINDOW_WIDTH//2 - 100, WINDOW_HEIGHT//2 - 50, 200, 50, "开始游戏")
    quit_button = Button(WINDOW_WIDTH//2 - 100, WINDOW_HEIGHT//2 + 20, 200, 50, "退出游戏")
    
    while True:
        screen.fill(MENU_BG)
        
        # 绘制标题
        try:
            font = pygame.font.Font('arial.ttf', 80)  # 尝试加载更大的字体
        except:
            font = pygame.font.Font(None, 80)  # 如果没有 Arial 字体，使用默认字体
        
        # 绘制游戏标题
        title = font.render("贪吃蛇", True, WHITE)  # 使用白色字体
        title_rect = title.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//4))
        screen.blit(title, title_rect)
        
        # 绘制最高分
        high_score_font = pygame.font.Font(None, 36)
        high_score_text = high_score_font.render(f"历史最高分: {high_score}", True, WHITE)  # 使用白色字体
        high_score_rect = high_score_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2 + 100))
        screen.blit(high_score_text, high_score_rect)
        
        # 绘制所有按钮
        start_button.draw(screen)
        quit_button.draw(screen)
        
        pygame.display.flip()
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            if start_button.handle_event(event):
                return  # 开始游戏
            
            if quit_button.handle_event(event):
                pygame.quit()
                sys.exit()
                
        pygame.time.Clock().tick(60)

# 添加游戏难度设置
DIFFICULTY_SPEEDS = {
    "简单": 8,
    "普通": 12,
    "困难": 16
}

# 添加游戏难度选择界面
def show_difficulty_menu():
    buttons = []
    y_pos = WINDOW_HEIGHT//2 - 100
    for diff in DIFFICULTY_SPEEDS.keys():
        btn = Button(WINDOW_WIDTH//2 - 100, y_pos, 200, 50, diff)
        buttons.append(btn)
        y_pos += 70

    while True:
        screen.fill(MENU_BG)
        
        # 绘制标题
        font = pygame.font.Font(None, 64)
        title = font.render("选择难度", True, TITLE_COLOR)
        title_rect = title.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//4))
        screen.blit(title, title_rect)
        
        # 绘制按钮
        for btn in buttons:
            btn.draw(screen)
        
        pygame.display.flip()
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            for i, btn in enumerate(buttons):
                if btn.handle_event(event):
                    return list(DIFFICULTY_SPEEDS.keys())[i]
        
        pygame.time.Clock().tick(60)

# 添加暂停界面
def show_pause_menu():
    # 创建半透明遮罩
    pause_surface = pygame.Surface((WINDOW_WIDTH, WINDOW_HEIGHT), pygame.SRCALPHA)
    pause_surface.fill(PAUSE_BG)
    
    resume_button = Button(WINDOW_WIDTH//2 - 100, WINDOW_HEIGHT//2 - 30, 200, 50, "继续游戏")
    menu_button = Button(WINDOW_WIDTH//2 - 100, WINDOW_HEIGHT//2 + 30, 200, 50, "返回菜单")
    
    while True:
        screen.blit(pause_surface, (0, 0))
        
        font = pygame.font.Font(None, 64)
        title = font.render("游戏暂停", True, WHITE)
        title_rect = title.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//3))
        screen.blit(title, title_rect)
        
        resume_button.draw(screen)
        menu_button.draw(screen)
        
        pygame.display.flip()
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return "resume"
                    
            if resume_button.handle_event(event):
                return "resume"
            if menu_button.handle_event(event):
                return "menu"

# 修改主游戏循环
def main():
    while True:
        show_menu()
        difficulty = show_difficulty_menu()
        game_speed = DIFFICULTY_SPEEDS[difficulty]
        
        clock = pygame.time.Clock()
        snake = Snake()
        food = Food()
        paused = False
        
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_UP and snake.direction != DOWN:
                        snake.direction = UP
                    elif event.key == pygame.K_DOWN and snake.direction != UP:
                        snake.direction = DOWN
                    elif event.key == pygame.K_LEFT and snake.direction != RIGHT:
                        snake.direction = LEFT
                    elif event.key == pygame.K_RIGHT and snake.direction != LEFT:
                        snake.direction = RIGHT
                    elif event.key == pygame.K_ESCAPE:
                        action = show_pause_menu()
                        if action == "menu":
                            return
                        paused = False
                    elif event.key == pygame.K_SPACE:  # 空格键暂停
                        paused = not paused

            if paused:
                continue

            if not snake.update():
                break

            if snake.get_head_position() == food.position:
                points = 3 if food.is_special else 1
                snake.length += points
                snake.score += points
                food.randomize_position()
                while (food.position[0] < BLOCK_SIZE or 
                       food.position[0] >= WINDOW_WIDTH - BLOCK_SIZE or
                       food.position[1] < BLOCK_SIZE or 
                       food.position[1] >= WINDOW_HEIGHT - BLOCK_SIZE):
                    food.randomize_position()

            screen.fill(BACKGROUND)
            draw_grid()
            draw_border()
            snake.render()
            food.render()
            
            # 显示分数和难度
            font = pygame.font.Font(None, 36)
            score_text = font.render(f'得分: {snake.score}', True, SCORE_COLOR)
            diff_text = font.render(f'难度: {difficulty}', True, SCORE_COLOR)
            
            score_rect = score_text.get_rect(topleft=(20, 20))
            diff_rect = diff_text.get_rect(topleft=(20, 60))
            
            pygame.draw.rect(screen, BORDER_COLOR, 
                           (15, 15, score_rect.width + 10, score_rect.height + 10), 
                           border_radius=5)
            pygame.draw.rect(screen, BORDER_COLOR, 
                           (15, 55, diff_rect.width + 10, diff_rect.height + 10), 
                           border_radius=5)
                           
            screen.blit(score_text, score_rect)
            screen.blit(diff_text, diff_rect)
            
            pygame.display.update()
            clock.tick(game_speed)

if __name__ == '__main__':
    main() 